package main

import (
	"3nets.io/wireguardcommon"
	"crypto/rand"
	"encoding/base64"
	"errors"
	"fmt"
	"github.com/BurntSushi/toml"
	"log"
	"net"
	"os"
	"reflect"
	"strconv"
	"strings"
	"time"
)

func SocketClient(ip string, port int, message string) {
	var sleepSeconds = time.Second * 5
	var sleepSecondIncrements = time.Second * 5
	var maxSleepSeconds = time.Second * 50
	addr := strings.Join([]string{ip, strconv.Itoa(port)}, ":")
	conn, err := net.Dial("tcp", addr)
	for err != nil {
		var ss = sleepSeconds % maxSleepSeconds
		fmt.Printf("could not connect.Sleeping for %s\n", ss)
		time.Sleep(ss)
		sleepSeconds += sleepSecondIncrements
		conn, err = net.Dial("tcp", addr)
	}

	//Keeping this code as a reference to the ResolveHostIP
	//log.Printf("connected on %s\n", conn.LocalAddr())
	//ipaddr := conn.LocalAddr().(*net.TCPAddr)
	//fmt.Printf("got ipaddr = %s\n", ipaddr)
	//_ = wireguardcommon.ResolveHostIp()

	defer func(conn net.Conn) {
		err := conn.Close()
		if err != nil {
			fmt.Println("error closing edge client connection to " + ip + ":" + string(rune(port)))
		}
	}(conn)

	_, err = conn.Write([]byte(message))
	if err != nil {
		log.Printf("failed to send message: %s\n", message)
		return
	}
	log.Printf("Send: %s", message)

	managerDHPublic := make([]byte, 1024)
	n, _ := conn.Read(managerDHPublic)
	managerDHPublicStr := string(managerDHPublic[:n])
	log.Printf("managerDHPublic: %s, n=%d", managerDHPublic[:n], n)

	presharedkey, edgeDHPublic, err := wireguardcommon.GenPskOnEdge(string(managerDHPublic[:n]))
	if err != nil {
		fmt.Printf("Failed to generate edge presharedkey: %s\n", err)
	} else {
		fmt.Printf("edgeDHPublic = %s\n", edgeDHPublic)
		fmt.Printf("presharedkey = %s\n", presharedkey)
	}

	_, err = conn.Write([]byte(edgeDHPublic))
	if err != nil {
		fmt.Printf("Failed to send edgeDHpublic, err: %s\n", err)
		return
	}

	encryptedAck := make([]byte, 1024)
	n, _ = conn.Read(encryptedAck)
	log.Printf("encryptedAck: %x, n=%d\n", encryptedAck[:n], n)
	secret, _ := base64.StdEncoding.DecodeString(presharedkey)
	ack, _ := wireguardcommon.Chacha20poly1305Decrypt(secret, encryptedAck[:n])
	log.Printf("ack: %x\n", ack)
	tokens := strings.Split(string(ack), ":")
	var endpointIP string
	var myWgIpAndMask string
	if len(tokens) >= 2 && tokens[0] == "OK" {
		endpointIP = tokens[1]
	}

	if len(tokens) >= 3 && tokens[0] == "OK" {
		myWgIpAndMask = tokens[2]
	}
	fmt.Printf("got OK and endpointIP %s\n", endpointIP)
	fmt.Printf("my ip and mask %s\n", myWgIpAndMask)

	///////////////////////////////
	// derive edge keys         //
	/////////////////////////////
	edgeWgPrivateStr, _, err := wireguardcommon.CreateKeyPairFromSeed(message)
	if err != nil {
		fmt.Printf("Failed to generate manager private/public key pair: %s\n", err)
		return
	}
	///////////////////////////////
	// end of derive edge keys  //
	/////////////////////////////
	edgeWgPrivateKey, err := base64.StdEncoding.DecodeString(edgeWgPrivateStr)
	if err != nil {
		fmt.Printf("failed to parse edgeWgPrivateStr: %s, err  %s\n", edgeWgPrivateStr, err)
		return
	}
	_, wireguardLink, err := wireguardcommon.ProvisionWireguardInterface(edgeWgPrivateKey, myWgIpAndMask)
	if err != nil {
		fmt.Printf("failed to create wireguard interface %s\n", err)
		return
	}

	managerIp, _ := wireguardcommon.GetFirstAddr()
	managerIpStr := managerIp.String() + "/32"
	fmt.Printf("adding peer ..\n")
	wireguardcommon.AddPeer(presharedkey, managerDHPublicStr, endpointIP, managerIpStr, true)
	fmt.Printf("adding linkAddr ..\n")
	err = wireguardcommon.AddLinkAddress(wireguardLink, myWgIpAndMask, managerIp.String())
	if err != nil {
		fmt.Printf("failed at AddLinkAddress, err: %s\n", err)
		return
	}

	err = wireguardcommon.StoreWgConfig()
	if err != nil {
		fmt.Printf("Failed to save wg conf, err: %s\n", err)
		os.Exit(wireguardcommon.ExitSetupFailed)
	}

	return
}

var edgeSignature_3nets_edge_io = "BOI8PrXNaVANosmOgylvmCMpL7Ek3AEAmvkrAJJ6fKxoJQTh1w0Tksv3uvXAVsxmfXeKE97il70Q8blrvu5rAjNuZXRzZWRnZS5pbw=="
var edgeSignature_edge1_3netsedge_io = "V4iNldDcxqLcCtb3vf254R0jlnNliS8G+ZJx9mkkikP5Is/XDdWw0PPoObebLdR9jcKrnjmRrPIO/D2kgaLNyrC9"

func main() {
	instanceFile := "/var/run/cloud-init/.instance-id"
	var conf wireguardcommon.ClientConfig

	var edgeSignature = edgeSignature_edge1_3netsedge_io
	conf.Signature = edgeSignature
	conf.InitialPort = wireguardcommon.InitialPort
	conf.ManagerIP = wireguardcommon.Ip2

	val := reflect.Indirect(reflect.ValueOf(conf))
	confSignature := val.Type().Field(0).Name
	confInitialPort := val.Type().Field(1).Name
	confManagerIp := val.Type().Field(2).Name

	//Try to read the configfile
	if _, err := os.Stat(wireguardcommon.ClientConfigFile); err == nil {
		fmt.Printf("found configfile at = %s\n", wireguardcommon.ClientConfigFile)
		var meta toml.MetaData
		conf, meta, err = wireguardcommon.ReadClientConfig(wireguardcommon.ClientConfigFile)
		if !meta.IsDefined(confSignature) {
			conf.Signature = edgeSignature
		}
		if !meta.IsDefined(confInitialPort) {
			conf.InitialPort = wireguardcommon.InitialPort
		}
		if !meta.IsDefined(confManagerIp) {
			conf.ManagerIP = wireguardcommon.Ip2
		}

	} else if errors.Is(err, os.ErrNotExist) {
		if _, err := os.Stat(instanceFile); err == nil {
			fmt.Printf("found instanceFile at = %s\n", instanceFile)
			var instanceId []byte
			instanceId, err = os.ReadFile(instanceFile)
			if err != nil {
			} else {
				conf.Signature += ":" + string(instanceId[:])
			}
		} else {
			//generate a nonce
			fmt.Printf("generating a nonce = %s\n", instanceFile)
			nonce := make([]byte, 32)
			if _, err := rand.Read(nonce); err != nil {
				return
			}
			nonceStr := base64.StdEncoding.EncodeToString(nonce)
			fmt.Printf("nonce = %s\n", nonceStr)

			//add the nonce to the edgeSignature
			conf.Signature += ":" + nonceStr
		}
	} else {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(wireguardcommon.ExitSetupFailed)
	}
	fmt.Printf("Signature = %s\n", conf.Signature)
	fmt.Printf("InitialPort = %d\n", conf.InitialPort)
	fmt.Printf("ManagerIp = %s\n", conf.ManagerIP)
	var err error
	if _, err = os.Stat(wireguardcommon.ConfFilename); err == nil {
		fmt.Println("found a saved config")
		_, wireguardLink, err := wireguardcommon.CreateWireguardInterface()
		if err != nil {
			fmt.Printf("Could not create the wireguardlink. err: %s\n", err)
		} else {
			specialLines, err := wireguardcommon.ConfigureWgFromSavedConfig()
			if err != nil {
				fmt.Println("Error loading from config, err:", err.Error())
			}
			fmt.Printf("special Lines thet were removed before applying\n%s\n", specialLines)
			for _, s := range specialLines {
				if strings.HasPrefix(s, "Address") {
					s = strings.Replace(s, "Address=", "", -1)
					s = strings.TrimSpace(s)
					exists, _ := wireguardcommon.HasLinkAddress(wireguardLink, s)
					if exists {
						fmt.Printf("LinkAddress (%s,%s) exists\n", wireguardLink.Type(), s)
					} else {
						err1 := wireguardcommon.AddLinkAddress(wireguardLink, s, "")
						if err1 != nil {
							//ignore err1. just print a message
							fmt.Printf("AddLinkAddress(%s,%s) failed, err: %s\n", wireguardLink.Type(), s, err)
						}
					}
				}
			}
		}
	}

	//Try to recreate it by contacting manager
	if err != nil {
		SocketClient(conf.ManagerIP, conf.InitialPort, conf.Signature)
	}
}
